print("开始对所有图片进行处理！")
import os
import fitz  # PyMuPDF
from PIL import Image, ImageDraw

# 设置PDF和图片文件夹路径
pdf_folder = "1.需要修改的PDF"
output_folder = "2.处理好的图"
if not os.path.exists(output_folder):
	os.makedirs(output_folder)

# 预定义的regions列表，每个元素包含适用于特定图片尺寸的矩形坐标
regions_by_size = {
	(1920, 1080): [
		((100, 100), (300, 200)),  # 第一个矩形
		((400, 100), (600, 200)),  # 第二个矩形
	],
	# 可以添加更多尺寸和对应的regions
	# (其他宽度, 其他高度): [
	#     ((..., ...), (..., ...)),
	#     ((..., ...), (..., ...)),
	# ],
}

# 遍历PDF中的每一页
for filename in os.listdir(pdf_folder):
	if filename.lower().endswith('.pdf'):
		pdf_path = os.path.join(pdf_folder, filename)
		with fitz.open(pdf_path) as doc:
			for page_num in range(len(doc)):
				# 获取PDF页面并转换为图片
				page = doc[page_num]
				pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))

				# 获取图片尺寸
				width, height = pix.width, pix.height

				# 根据图片尺寸选择regions
				regions_to_use = regions_by_size.get((width, height), None)
				if regions_to_use is None:
					print(f"没有找到适用于 {width}x{height} 尺寸的regions，跳过处理...")
					continue

				# 将 PyMuPDF 的 pixmap 对象转换为 PIL Image 对象
				image = Image.frombuffer(
					'RGB',
					(pix.width, pix.height),
					pix.samples,
					'raw',
					'RGB',
					pix.stride,
					1
				)

				# 创建一个可以在图片上绘制的对象
				draw = ImageDraw.Draw(image)

				# 遍历每个区域并填充白色
				for region in regions_to_use:
					left_top = region[0]
					right_bottom = region[1]
					x1, y1 = left_top
					x2, y2 = right_bottom
					# 确保y2不小于y1，否则交换并调整绘制方向
					if y1 > y2:
						y1, y2 = y2, y1
					draw.rectangle((x1, y1, x2, y2), fill=(255, 255, 255))  # 填充白色

				# 保存修改后的图片
				output_image_path = os.path.join(
					output_folder,
					f"{filename.rsplit('.', 1)[0]}_page_{page_num + 1}.png"  # 添加页码以避免文件名冲突
				)
				image.save(output_image_path)

				# 打印图片处理完成的消息（如果需要）
				print(f"{output_image_path} 图片处理完成~")

print("已完成所有图片的处理。")